#include "kpextOpImgCMapF32C1toU8C4.h"

using namespace kpext;

#define INPUT_MAX_V 65535
#define COLOR_MAP_LEN 1024

typedef struct{
    float r;
    float g;
    float b;
    float a;
}FRGBA_SZ;

std::mutex OpImgCMapF32C1U8C4::m_CMapsMutex;

std::map<uint32_t, std::shared_ptr<float>> OpImgCMapF32C1U8C4::m_ColorMaps({});

OpImgCMapF32C1U8C4::OpImgCMapF32C1U8C4(kpextVkContext *ctx) :
    BaseCompute(ctx)
{
    if(!m_vkCtx){
        throw std::runtime_error(__FUNCTION__ " construct with null pointer");
    }

    m_cols = 1;
    m_rows = 1;
    m_cmap_idx = ~0U;
    m_inBuf = nullptr;
    m_outBuf = nullptr;
    m_cmapBuf = nullptr;

    m_postBarrier = true;

    {
        //  颜色映射初始化
        std::lock_guard<std::mutex> lk(m_CMapsMutex);
        if(m_ColorMaps.find(0) == m_ColorMaps.end()
                || m_ColorMaps[0] == nullptr){
            std::shared_ptr<float> cm (new float[COLOR_MAP_LEN * 4 * sizeof (float)]);
            auto ptr = reinterpret_cast<FRGBA_SZ *>(cm.get());
            for (int i = 0; i < COLOR_MAP_LEN; ++i) {
                float v = static_cast<float>(i) / 1023.0f;
                v = v * 255.0f;
                ptr[i].r = v;
                ptr[i].g = v;
                ptr[i].b = v;
                ptr[i].a = 255.0f;
            }
            m_ColorMaps[0] = cm;
        }
    }

    //  初始化描述符集及布局
    m_descrPool = m_vkCtx->descriptorPool(
                {2, 1}, //{存储缓冲,uniform纹素缓冲,uniform缓冲,纹素缓冲,存储图像}
                1);

    vk::DescriptorSetLayoutBinding binding(
                0,
                vk::DescriptorType::eStorageBuffer,
                1,
                vk::ShaderStageFlagBits::eCompute
                );
    std::vector<vk::DescriptorSetLayoutBinding> bindings;
    bindings.push_back(
                binding
                );

    binding.setBinding(1);
    bindings.push_back(binding);

    binding.setBinding(2);
    binding.setDescriptorType(
                vk::DescriptorType::eUniformTexelBuffer);
    bindings.push_back(binding);

    m_descrSetlayout = m_descrPool->createDescriptorSetLayout(
                bindings
                );

    m_descrSet = m_descrPool->recreateDescSet(*m_descrSetlayout,
                                              m_descrSet
                                              );

    // 初始化pipeline
    m_pipeline = m_vkCtx->computePipeline();
    m_pipeline->m_specConsts = std::make_shared<PipelineConstSZ>(5, 4);
    m_pipeline->m_pushConsts = std::make_shared<PipelineConstSZ>(2, 4);

    m_workGroupSize = ctx->getDefWorkGroupSize();

    m_pipeline->m_specConsts->set({m_workGroupSize[0],
                                   m_workGroupSize[1],
                                   1
                                  });
    m_pipeline->m_specConsts->itemFloat(3) = INPUT_MAX_V;
    m_pipeline->m_specConsts->itemUint(4) = COLOR_MAP_LEN;

    m_pipeline->m_pushConsts->set({1, 1});

    m_pipeline->m_descrSetLayouts.push_back(m_descrSetlayout);

    if(!m_pipeline->rebuild(
                sizeof (spirv::ImgCMapF32C1toU8C4), spirv::ImgCMapF32C1toU8C4)
            ){
        SPDLOG_WARN("Rebuild kpext::ComputePipeline failed !");
    }

    //  工作组数量
    m_workCount = calc2DWorkGroup(1, 1,
                                  m_workGroupSize);
}

OpImgCMapF32C1U8C4::~OpImgCMapF32C1U8C4()
{

}

std::shared_ptr<float> OpImgCMapF32C1U8C4::getColorMapClone(uint32_t idx)
{
    std::lock_guard<std::mutex> lk(m_CMapsMutex);
    if(m_ColorMaps.find(idx) == m_ColorMaps.end()
            ||
            ! m_ColorMaps[idx]
            ){
        return nullptr;
    }

    size_t size = COLOR_MAP_LEN * 4 * sizeof (float);
    std::shared_ptr<float> ret(new float[size]);
    std::memcpy(ret.get(), m_ColorMaps[idx].get(), size);
    return ret;
}

void OpImgCMapF32C1U8C4::setColorMap(uint32_t idx,
                                     const std::vector<float> &mapdata)
{
    std::lock_guard<std::mutex> lk(m_CMapsMutex);
    auto size = mapdata.size();
    if(size < 10 || idx == ~0U){
        SPDLOG_ERROR(__FUNCTION__ " invalid input !");
        return;
    }

    int count = static_cast<int>(size / 5);
    std::map<float, std::array<float, 4>> tempmap;

    for (int i = 0; i < count; ++i) {
        size_t idx = static_cast<size_t>(i * 5);
        tempmap[mapdata[idx]] = {mapdata[idx + 1],
                                 mapdata[idx + 2],
                                 mapdata[idx + 3],
                                 mapdata[idx + 4]};
    }


    size = COLOR_MAP_LEN * 4 * sizeof (float);
    std::shared_ptr<float> cm (new float[size]);
    std::memset(cm.get(), 0, size);

    auto it = tempmap.end();
    float scale = (--it)->first;
    float min = tempmap.begin()->first;
    scale = scale - min;

    count = static_cast<int>(tempmap.size()) - 1;
    for (int i = 0; i < COLOR_MAP_LEN; ++i) {
        float v = static_cast<float>(i) / 1023.0f;
        v = scale * v + min;

        it = tempmap.begin();
        auto it1 = tempmap.begin();
        it1++;
        for (int j = 0; j < count; ++j) {
            if(v >= it->first && v <= it1->first){
                float * dst = cm.get() + (i * 4);
                v = v - it->first;
                v = v / (it1->first - it->first);
                //                std::cout << v << std::endl;
                for (size_t c = 0; c < 4; ++c) {
                    dst[c] = v * (it1->second[c] - it->second[c])
                            + it->second[c];
                }
                break;
            }
            it++;
            it1++;
        }
    }

    m_ColorMaps[idx] = cm;
}


void OpImgCMapF32C1U8C4::rebuild(int cols, int rows,
                                 uint32_t cmap,
                                 std::shared_ptr<BaseBuffer> inbuf,
                                 std::shared_ptr<BaseBuffer> outbuf,
                                 std::shared_ptr<CommandBuffer> cmdbuf)
{
    if(!inbuf || !outbuf || !cmdbuf || !m_vkCtx){
        throw std::runtime_error(__FUNCTION__ " input with null pointer!");
    }

    if(cols < 0) cols = 1;
    if(rows < 0) rows = 1;

    m_cmdBarrier.prepare(
                {outbuf,},
                vk::PipelineStageFlagBits::eComputeShader,
                vk::PipelineStageFlagBits::eComputeShader,
                vk::AccessFlagBits::eShaderWrite,
                vk::AccessFlagBits::eShaderRead
                );

    //  颜色映射数据刷新
    if(m_cmapBuf == nullptr){
        m_cmapBuf = m_vkCtx->uniformTexelBuf(COLOR_MAP_LEN * 4 * sizeof (float),
                                             vk::Format::eR32G32B32A32Sfloat
                                             );
    }

    {
        std::lock_guard<std::mutex> lk(m_CMapsMutex);
        if(m_ColorMaps.find(cmap) == m_ColorMaps.end()){
            cmap = 0;
        }
    }

    if(m_cmap_idx != cmap || m_cmap_idx == ~0UL){
        m_cmap_idx = cmap;
        refreshColorMap(cmdbuf);
    }

    STCAST_ASSIGN(m_cols, cols);
    STCAST_ASSIGN(m_rows, rows);

    if(m_inBuf != inbuf.get()
            || m_outBuf != outbuf.get()
            ){
        BaseCompute::rebuild({inbuf, outbuf, m_cmapBuf});
        m_inBuf = inbuf.get();
        m_outBuf = outbuf.get();
    }

    //  设定输入输出尺寸
    m_pipeline->m_pushConsts->set(
                {m_cols, m_rows,
                }
                );

    //  重新计算工作组大小
    m_workCount = calc2DWorkGroup(m_cols, m_rows,
                                  m_workGroupSize);
}

void OpImgCMapF32C1U8C4::refreshColorMap(std::shared_ptr<CommandBuffer> cmdbuf)
{
    auto size = COLOR_MAP_LEN * 4 * sizeof (float);
    auto stagebuf = m_vkCtx->stagingBuf(size);
    {
        std::lock_guard<std::mutex> lk(m_CMapsMutex);
        std::memcpy(stagebuf->map(),
                    m_ColorMaps[m_cmap_idx].get(),
                    size
                    );
        stagebuf->unmap();
    }

    cmdbuf = cmdbuf->allocateCmdBuf(true,
                                    vk::CommandPoolCreateFlagBits::eTransient
                                    );
    kpext::CmdBaseBufferCpy bufcpy;
    bufcpy.prepare(
                m_cmapBuf,
                stagebuf,
                size
                );
    kpext::CmdBuffersMemBarrier barrierTx;
    barrierTx.prepare(
                m_cmapBuf,
                vk::PipelineStageFlagBits::eTransfer,
                vk::PipelineStageFlagBits::eComputeShader,
                vk::AccessFlagBits::eTransferWrite,
                vk::AccessFlagBits::eShaderRead
                );
    cmdbuf->begin();
    cmdbuf->record(&bufcpy);
    cmdbuf->record(&barrierTx);
    cmdbuf->end();
    cmdbuf->submit();
    cmdbuf->waitExec();
}
